import torch
import torch.nn as nn
import torch.optim as optim
from game.game import SnakeGame
from model.replay_buffer import ReplayBuffer
from train import train


class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(14, 64)
        self.fc2 = nn.Linear(64, 4)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.softmax(self.fc2(x))


def main_simple_nn():
    env = SnakeGame()
    simple_nn_model = SimpleNN()
    buffer = ReplayBuffer()
    optimizer = optim.AdamW(simple_nn_model.parameters(), lr=1e-4, weight_decay=1e-5)

    # 训练过程
    train(simple_nn_model, env, buffer, optimizer, episodes=2000)


if __name__ == "__main__":
    main_simple_nn()
